package javaagent.mybatis;

import javassist.*;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.security.ProtectionDomain;
import java.util.Date;

/**
 * 使用反射方案解决双亲委派的限制问题
 */
public class MyBatisAgent2 implements ClassFileTransformer {

    @Override
    public byte[] transform(ClassLoader loader, String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) {
        if (!"org/apache/ibatis/executor/BaseExecutor".equals(className)) {
            return null;
        }
        try {
            ClassPool pool = new ClassPool();
            pool.appendSystemPath();
            pool.appendClassPath(new LoaderClassPath(loader));
            CtClass ctClass = pool.get("org.apache.ibatis.executor.BaseExecutor");
            // 查询
            CtMethod ctMethod = ctClass.getDeclaredMethods("query")[1];
            ctMethod.addLocalVariable("info", pool.get(SqlInfo.class.getName()));
            ctMethod.insertBefore("info=javaagent.mybatis.MyBatisAgent2.begin($args);");
            ctMethod.insertAfter("javaagent.mybatis.MyBatisAgent2.end(info);");
            System.out.println("插桩成功：" + ctClass.getName());
            return ctClass.toBytecode();
        } catch (NotFoundException | CannotCompileException | IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static SqlInfo begin(Object[] params) {
        SqlInfo info = new SqlInfo();
        info.beginTime = System.currentTimeMillis();
        BoundSqlAdapter adapter = new BoundSqlAdapter(params[5]);
        info.sql = adapter.getSql();
        return info;
    }

    public static void end(SqlInfo info) {
        info.useTime = System.currentTimeMillis() - info.beginTime;
        System.out.println(info);
    }

    public static class SqlInfo {
        public long beginTime;
        public long useTime;
        public String sql;

        @Override
        public String toString() {
            return "SqlInfo{" +
                    "beginTime=" + new Date(beginTime) +
                    ", useTime=" + useTime +
                    ", sql='" + sql + '\'' +
                    '}';
        }
    }

    /**
     * 通过反射获取BoundSQL
     */
    public static class BoundSqlAdapter {
        Object target;
        private static Method getSql;
        private static Class aClass;

        private synchronized static void init(Class cls) {
            try {
                aClass = cls;
                getSql = cls.getDeclaredMethod("getSql");
                getSql.setAccessible(true);
            } catch (NoSuchMethodException e) {
                throw new RuntimeException(e);
            }

        }

        public BoundSqlAdapter(Object target) {
            this.target = target;
            if (aClass == null) {
                init(target.getClass());
            }
            this.target = target;
        }

        public String getSql() {
            try {
                return (String) getSql.invoke(target);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }
}
